/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cmath>
#include <iostream>
#include "vec_dot.h"
#include "parvec_axpby.h"
#include "parvec_axpbypcz.h"
#include "par_pipebicgstab_solver.h"

namespace YHAMG
{

Par_PipeBiCGStab_Solver::Par_PipeBiCGStab_Solver(const Par_Operator& _A)
	: A(&_A),
	P(0),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeBiCGStab_Solver::Par_PipeBiCGStab_Solver(const Par_Operator& _A, const Par_Operator& _P)
	: A(&_A),
	P(&_P),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeBiCGStab_Solver::Par_PipeBiCGStab_Solver(const Par_Operator& _A, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(0),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

Par_PipeBiCGStab_Solver::Par_PipeBiCGStab_Solver(const Par_Operator& _A, const Par_Operator& _P, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(&_P),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

#define PBCGS_RR_STEPS         100
#define APPLYDOT(x, y, result) ApplyDot(x, y, result, buf, dst, cnt)
#define BEGINREDUCTION()       BeginReduction(buf, cnt, comm, &request)
#define ENDREDUCTION()         EndReduction(buf, dst, cnt, &request, &status)

static inline void ApplyDot(const Par_Vector& x, const Par_Vector& y, double& result, double* buf, double **dst, int &cnt)
{
	buf[cnt] = Vec_Dot(x.Loc, y.Loc);
	dst[cnt++] = &result;
}

static inline void BeginReduction(double* buf, int cnt, MPI_Comm comm, MPI_Request* request)
{
	MPI_Iallreduce(MPI_IN_PLACE, buf, cnt, MPI_DOUBLE, MPI_SUM, comm, request);
}

static inline void EndReduction(const double* buf, double** dst, int &cnt, MPI_Request* request, MPI_Status* status)
{
	MPI_Wait(request, status);
	for (int i = 0; i < cnt; ++i)
		*dst[i] = buf[i];
	cnt = 0;
}

void Par_PipeBiCGStab_Solver::operator()(const Par_Vector& b, const Par_Vector& x, int& iter, double& relres) const
{
	MPI_Comm comm = A->comm;

	int rank;
	MPI_Comm_rank(comm, &rank);

	MPI_Request request;
	MPI_Status status;
	double buf[5];
	double* dst[5];
	int cnt = 0;

	int n = A->OutSize();

	double normb, res, alpha, beta, omega, rho, rho1, qy, yy, r0w, r0s, r0z;
	Par_Vector r(comm), r0(comm), r1(comm), s(comm), s1(comm), t(comm), p1(comm), q(comm), q1(comm), v(comm), w(comm), w1(comm), y(comm), z(comm), z1(comm);

	r.Resize(n);
	r0.Resize(n);
	r1.Resize(n);
	s.Resize(n);
	s1.Resize(n);
	t.Resize(n);
	p1.Resize(n);
	q.Resize(n);
	q1.Resize(n);
	v.Resize(n);
	w.Resize(n);
	w1.Resize(n);
	y.Resize(n);
	z.Resize(n);
	z1.Resize(n);

	if (PrintStats && rank == 0)
	{
		std::cout << "Iter\tResidual\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	iter = 0;

	APPLYDOT(b, b, normb);

	BEGINREDUCTION();

	A->Apply(x, r);
	ParVec_AXPBY(-1.0, r, 1.0, b, r);

	r0.Copy(r);

	if (P)
		P->Apply(r, r1);
	else
		r1.Copy(r);

	A->Apply(r1, w);

	ENDREDUCTION();

	normb = sqrt(normb);

	while (1)
	{
		APPLYDOT(r0, r, rho);
		APPLYDOT(r0, w, r0w);

		if (iter > 0)
		{
			APPLYDOT(r0, s, r0s);
			APPLYDOT(r0, z, r0z);
			APPLYDOT(r, r, res);
		}

		BEGINREDUCTION();

		if (P)
			P->Apply(w, w1);
		else
			w1.Copy(w);

		A->Apply(w1, t);

		ENDREDUCTION();

		if (iter == 0) res = rho;

		res = sqrt(res);

		if (PrintStats && rank == 0)
			std::cout << iter << "\t" << res << "\n";

		if (res / normb <= Tolerance || iter == MaxIters) break;

		if (iter == 0)
		{
			alpha = rho / r0w;

			p1.Copy(r1);
			s.Copy(w);
			s1.Copy(w1);
			z.Copy(t);
		}
		else
		{
			beta = (alpha / omega) * (rho / rho1);
			alpha = rho / (r0w + beta * r0s - beta * omega * r0z);

			ParVec_AXPBYPCZ(1.0, r1, beta, p1, -beta * omega, s1, p1);
			ParVec_AXPBYPCZ(1.0, w, beta, s, -beta * omega, z, s);
			ParVec_AXPBYPCZ(1.0, w1, beta, s1, -beta * omega, z1, s1);
			ParVec_AXPBYPCZ(1.0, t, beta, z, -beta * omega, v, z);
		}

		ParVec_AXPBY(1.0, r, -alpha, s, q);
		ParVec_AXPBY(1.0, r1, -alpha, s1, q1);
		ParVec_AXPBY(1.0, w, -alpha, z, y);

		APPLYDOT(q, y, qy);
		APPLYDOT(y, y, yy);

		BEGINREDUCTION();

		if (P)
			P->Apply(z, z1);
		else
			z1.Copy(z);

		A->Apply(z1, v);

		ENDREDUCTION();

		omega = qy / yy;

		x.Add2Scaled(alpha, p1, omega, q1);

		ParVec_AXPBY(1.0, q, -omega, y, r);
		ParVec_AXPBYPCZ(1.0, q1, -omega, w1, omega * alpha, z1, r1);
		ParVec_AXPBYPCZ(1.0, y, -omega, t, omega * alpha, v, w);

		rho1 = rho;

		++iter;
		
		if (iter % PBCGS_RR_STEPS == 0)
		{
			A->Apply(x, r);
			ParVec_AXPBY(-1.0, r, 1.0, b, r);
			if (P)
				P->Apply(r, r1);
			else
				r1.Copy(r);
			A->Apply(r1, w);
			A->Apply(p1, s);
			if (P)
				P->Apply(s, s1);
			else
				s1.Copy(s);
			A->Apply(s1, z);
		}
	}

	relres = res / normb;

	if (PrintStats && rank == 0)
	{
		std::cout << "----------------------------------------------------------------\n";
		std::cout << "Iterations: " << iter << "\n";
		std::cout << "Final Relative Residual: " << relres << "\n";
	}
}

}